import numpy as np
import xgboost as xgb
from matplotlib import cm
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


fig = plt.figure()
ax2 = Axes3D(fig)
divide_num = 2.0
xx = np.arange(-10, 10, divide_num)
yy = np.arange(-10, 10, divide_num)
X, Y = np.meshgrid(xx, yy)
Z = 1 - X - Y
print(X.shape)
print()
print(Y.shape)
print()
print(Z.shape)
print("==================")
total_length = int(20.0 / divide_num * 20.0 / divide_num)
print(X.reshape(total_length))
X_train = np.vstack((X.reshape(total_length), Y.reshape(total_length))).T
y_train = Z.reshape(total_length)
print(X_train)
print(y_train)

# Normal Regressor Method
model = xgb.XGBRegressor()
model.fit(X_train, y_train)
y_predict = model.predict(X_train)
y_predict = y_predict.reshape(Z.shape)

# Special Method


ax2.plot_wireframe(X, Y, Z, rstride=1, cstride=1)
ax2.plot_surface(X, Y, y_predict, cmap=cm.coolwarm)
plt.show()
plt.close()
